import torch
import torch.nn as nn
import torch.nn.functional as f
from torch.nn import init


class BNConvReLU(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):  #, activation='relu', norm='BN'):
        super(BNConvReLU, self).__init__()

        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False)
        #self.activation = getattr(torch, 'relu', 'relu')
        self.activation = nn.ReLU(inplace=False)
        self.norm_layer = nn.BatchNorm2d(out_channels)


    def forward(self, x):
        out = self.conv2d(x)
        out = self.norm_layer(out)
        out = self.activation(out)

        return out



class BNConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):  #, norm='BN'):
        super(BNConv, self).__init__()

        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False)
        self.norm_layer = nn.BatchNorm2d(out_channels)


    def forward(self, x):
        out = self.conv2d(x)
        out = self.norm_layer(out)

        return out


class ConvReLU(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        super(ConvReLU, self).__init__()

        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=True)
        self.activation = nn.ReLU(inplace=False)

    def forward(self, x):
        out = self.conv2d(x)
        out = self.activation(out)

        return out


class SELayer(nn.Module):
    def __init__(self, channel, reduction=16):
        super(SELayer, self).__init__()


        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=False),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)





class ConvLSTM(nn.Module):
    """Adapted from: https://github.com/Atcold/pytorch-CortexNet/blob/master/model/ConvLSTMCell.py """

    def __init__(self, input_size, hidden_size, kernel_size):
        super(ConvLSTM, self).__init__()

        self.conv = ConvReLU(input_size, hidden_size, kernel_size=3, stride=2, padding=1)

        self.input_size = hidden_size
        self.hidden_size = hidden_size
        pad = kernel_size // 2

        # cache a tensor filled with zeros to avoid reallocating memory at each inference step if --no-recurrent is enabled
        self.zero_tensors = {}

        self.Gates = nn.Conv2d(hidden_size + hidden_size, 4 * hidden_size, kernel_size, padding=pad)


    def forward(self, input_, prev_state=None):

        input_ = self.conv(input_)

        # get batch and spatial sizes
        batch_size = input_.data.size()[0]
        spatial_size = input_.data.size()[2:]

        # generate empty prev_state, if None is provided
        if prev_state is None:
            # create the zero tensor if it has not been created already
            state_size = tuple([batch_size, self.hidden_size] + list(spatial_size))
            if state_size not in self.zero_tensors:
                # allocate a tensor with size `spatial_size`, filled with zero (if it has not been allocated already)
                self.zero_tensors[state_size] = (
                    torch.zeros(state_size).to(input_.device),
                    torch.zeros(state_size).to(input_.device)
                )
            prev_state = self.zero_tensors[tuple(state_size)]
        prev_hidden, prev_cell = prev_state

        # data size is [batch, channel, height, width]
        stacked_inputs = torch.cat((input_, prev_hidden), 1)
        gates = self.Gates(stacked_inputs)

        # chunk across channel dimension
        in_gate, remember_gate, out_gate, cell_gate = gates.chunk(4, 1)

        # apply sigmoid non linearity
        in_gate2 = torch.sigmoid(in_gate)
        remember_gate2 = torch.sigmoid(remember_gate)
        out_gate2 = torch.sigmoid(out_gate)

        # apply tanh non linearity
        cell_gate2 = torch.tanh(cell_gate)

        # compute current cell and hidden state
        cell = (remember_gate2 * prev_cell) + (in_gate2 * cell_gate2)
        hidden = out_gate2 * torch.tanh(cell)

        return hidden, cell

class ConvLSTM2(nn.Module):
    """Adapted from: https://github.com/Atcold/pytorch-CortexNet/blob/master/model/ConvLSTMCell.py """

    def __init__(self, input_size, hidden_size, kernel_size):
        super(ConvLSTM2, self).__init__()

        self.conv = ConvReLU(input_size, hidden_size, kernel_size=kernel_size, stride=2, padding=1)
        self.rnn = nn.LSTMCell(hidden_size, hidden_size, num_layers=1, batch_first=False)

    def forward(self, input_):

        input_2 = self.conv(input_)
        #print(input_2.size())
        input_3 = input_2.view(-1, input_2.size()[2], input_2.size()[3])
        #print(input_3.size())
        out, (h, c) = self.rnn(input_3)
        x = out[:, -1, :]
        x.view(32, -1, x.size()[1], x.size()[2])

        return x